<?php
/**
 * Programmer: Owen Zhang
 * Email: owen@owenzhang.com
 * Start Date: 05/15/22
 * Last Update: 05 15, 2022 [OZ]
 * Functions:
 * 基于 Lua 脚本简单限流令牌桶，应用接口服务限流工具
 * 在原作者的基础上做的修改
 * 原地址：
 *      https://www.workerman.net/plugin/37
 *      https://github.com/Tinywan/webman-limit-traffic
 */

namespace app\middleware;

use app\common\ApiStatus;
use Webman\MiddlewareInterface;
use Webman\Http\Response;
use Webman\Http\Request;
use support\Redis;

class LimitTraffic implements MiddlewareInterface
{
    public const LIMIT_TRAFFIC_SCRIPT_SHA = 'limit:traffic:script';
    public const LIMIT_TRAFFIC_PRE        = 'limit:traffic:pre:';

    public function process(Request $request, callable $next): Response
    {
        if ($result = $this->traffic()) {
            return new Response($result['status'], [
                'Content-Type' => 'application/json',
            ], json_encode($result['body']));
        }
        return $next($request);
    }

    /**
     * @desc   : 返回允许的请求的最大数目及时间，例如，[100, 2] 表示在 2 秒内最多 100 次的 API 调
     * @author Tinywan(ShaoBo Wan)
     */
    public function getRateLimit(): array
    {
        return [
            'limit'       => 100, // 请求次数
            'window_time' => 1, // 窗口时间，单位：秒
            'status'      => 429,  // HTTP 状态码
            'body'        => [
                // 响应返回的结构信息
                'code' => ApiStatus::API_ERROR->value,
                'msg'  => 'Too Many Requests'
            ],
        ];
    }

    /**
     * 校测
     *
     * @return array|false
     */
    public function traffic()
    {
        $config    = $this->getRateLimit();
        $scriptSha = Redis::get(self::LIMIT_TRAFFIC_SCRIPT_SHA);
        if (!$scriptSha) {
            $script    = <<<luascript
            local result = redis.call('SETNX', KEYS[1], 1);
            if result == 1 then
                return redis.call('expire', KEYS[1], ARGV[2])
            else
                if tonumber(redis.call("GET", KEYS[1])) >= tonumber(ARGV[1]) then
                    return 0
                else
                    return redis.call("INCR", KEYS[1])
                end
            end
luascript;
            $scriptSha = Redis::script('load', $script);
            Redis::set(self::LIMIT_TRAFFIC_SCRIPT_SHA, $scriptSha);
        }
        $limitKey = self::LIMIT_TRAFFIC_PRE . request()->getRealIp(true);
        $result   = Redis::rawCommand('evalsha', $scriptSha, 1, $limitKey, $config['limit'], $config['window_time']);
        if ($result === 0) {
            return [
                'limit'     => $config['limit'],
                'remaining' => $config['limit'] - Redis::get($limitKey),
                'reset'     => Redis::ttl($limitKey),
                'status'    => $config['status'],
                'body'      => $config['body'],
            ];
        }
        return false;
    }
}